#!/usr/bin/env python3
# Author: Armit
# Create Time: 周五 2025/08/22

# 使用预训练whisper模型进行编码 (Encode)，查看两个dev集样本的编码余弦相似度

import os
from pathlib import Path
import torch
import torch.nn.functional as F
import matplotlib.pyplot as plt
import librosa as L

from utils import DEV_PATH, load_dev_index_file, get_whisper_model, get_whisper_latent

# softmax temperature (越小越接近onehot)
T = 0.1


@torch.inference_mode
def run():
  model = 'base'
  model_path = f'openai/whisper-{model}'
  model, processor = get_whisper_model(model_path)
  HOP_LEN = processor.feature_extractor.hop_length    # 160

  items = load_dev_index_file()
  for anc_wav_fp, com_wav_fp, truth in items:
    # 魔改成降噪板本
    com_wav_fp = Path(str(com_wav_fp).replace('Dev', 'Dev.mdx').replace('.wav', '_vocals.wav'))
    # [N=1, L=100, D=512]
    anc_wav, _ = L.load(anc_wav_fp, sr=None)
    com_wav, _ = L.load(com_wav_fp, sr=None)
    if len(anc_wav) > len(com_wav):
      anc_wav, com_wav = com_wav, anc_wav
    nframe1 = len(anc_wav) // HOP_LEN
    nframe2 = len(com_wav) // HOP_LEN
    latent1 = get_whisper_latent(anc_wav, model, processor).last_hidden_state[0, :nframe1]
    latent2 = get_whisper_latent(com_wav, model, processor).last_hidden_state[0, :nframe2]

    latent1 = F.normalize(latent1, p=2, dim=-1)   # [L1, D]
    latent2 = F.normalize(latent2, p=2, dim=-1)   # [L2, D]
    simmat = latent1 @ latent2.T                  # [L1, L2]
    simmat1 = F.softmax(simmat / T, dim=1)
    simmat2 = F.softmax(simmat / T, dim=0)
    # diagonal strengths
    pred = round((torch.trace(simmat1).item() + torch.trace(simmat2).item()) / 2, 5)

    if not 'reweight':
      # 交给encoder下一层
      latent1_new = simmat1   @ latent2
      latent2_new = simmat2.T @ latent1
      latent1 = latent1_new
      latent2 = latent2_new

    plt.clf()
    plt.subplot(121) ; plt.imshow(simmat1.cpu(), cmap='grey') ; plt.gca().invert_yaxis()
    plt.subplot(122) ; plt.imshow(simmat2.cpu(), cmap='grey') ; plt.gca().invert_yaxis()
    plt.suptitle(f'name: {Path(anc_wav_fp).stem}, pred: {pred}, truth: {truth}')
    plt.show()


if __name__ == '__main__':
  cwd = os.getcwd()
  os.chdir(DEV_PATH)
  try:
    run()
  finally:
    os.chdir(cwd)
